"""
内存管理器
负责执行引擎的内存分配、回收和监控
"""

import threading
from typing import Dict, Optional, Set
from collections import defaultdict
import gc
import psutil
import os

class MemoryPool:
    """内存池管理"""
    
    def __init__(self, pool_size: int = 64 * 1024 * 1024):  # 64MB默认
        self.pool_size = pool_size
        self.allocated_size = 0
        self.free_blocks: Dict[int, Set[int]] = defaultdict(set)  # 大小 -> 偏移集合
        self.allocated_blocks: Dict[int, int] = {}  # 偏移 -> 大小
        self.lock = threading.RLock()
    
    def allocate(self, size: int) -> Optional[int]:
        """分配内存块，返回偏移量"""
        with self.lock:
            # 找到合适大小的块
            actual_size = self._round_up_to_power_of_2(size)
            
            # 检查是否有可用的块
            if actual_size in self.free_blocks and self.free_blocks[actual_size]:
                offset = self.free_blocks[actual_size].pop()
                self.allocated_blocks[offset] = actual_size
                return offset
            
            # 检查是否有更大的块可以分割
            for available_size in sorted(self.free_blocks.keys()):
                if available_size > actual_size and self.free_blocks[available_size]:
                    offset = self.free_blocks[available_size].pop()
                    
                    # 分割块
                    self.allocated_blocks[offset] = actual_size
                    
                    # 将剩余部分放回空闲列表
                    remaining_size = available_size - actual_size
                    if remaining_size > 0:
                        self.free_blocks[remaining_size].add(offset + actual_size)
                    
                    return offset
            
            # 检查是否有足够空间分配新块
            if self.allocated_size + actual_size <= self.pool_size:
                offset = self.allocated_size
                self.allocated_size += actual_size
                self.allocated_blocks[offset] = actual_size
                return offset
            
            return None  # 内存不足
    
    def deallocate(self, offset: int):
        """释放内存块"""
        with self.lock:
            if offset not in self.allocated_blocks:
                return
            
            size = self.allocated_blocks.pop(offset)
            self.free_blocks[size].add(offset)
            
            # 尝试合并相邻的空闲块
            self._try_merge_blocks(offset, size)
    
    def _round_up_to_power_of_2(self, size: int) -> int:
        """将大小向上舍入到2的幂"""
        if size <= 0:
            return 1
        
        power = 1
        while power < size:
            power <<= 1
        return power
    
    def _try_merge_blocks(self, offset: int, size: int):
        """尝试合并相邻的空闲块"""
        # 简化的合并逻辑，实际实现会更复杂
        pass
    
    def get_usage_stats(self) -> Dict:
        """获取内存使用统计"""
        with self.lock:
            allocated_bytes = sum(self.allocated_blocks.values())
            free_bytes = self.pool_size - self.allocated_size
            
            return {
                'pool_size': self.pool_size,
                'allocated_bytes': allocated_bytes,
                'free_bytes': free_bytes,
                'utilization': allocated_bytes / self.pool_size if self.pool_size > 0 else 0,
                'fragmentation': len(self.free_blocks)
            }

class MemoryManager:
    """内存管理器主类"""
    
    def __init__(self, max_memory_mb: int = 512):
        self.max_memory_bytes = max_memory_mb * 1024 * 1024
        self.pools: Dict[str, MemoryPool] = {}
        self.pool_sizes = {
            'small': 16 * 1024 * 1024,    # 16MB
            'medium': 64 * 1024 * 1024,   # 64MB  
            'large': 256 * 1024 * 1024,   # 256MB
        }
        self.lock = threading.RLock()
        self.allocated_objects: Dict[int, str] = {}  # 对象ID -> 池名
        
        # 初始化内存池
        for pool_name, pool_size in self.pool_sizes.items():
            self.pools[pool_name] = MemoryPool(pool_size)
    
    def allocate(self, size: int, pool_name: str = None) -> Optional[int]:
        """分配内存"""
        with self.lock:
            if pool_name is None:
                pool_name = self._select_pool_for_size(size)
            
            if pool_name not in self.pools:
                return None
            
            # 检查内存限制
            if self._get_total_allocated() + size > self.max_memory_bytes:
                # 尝试垃圾回收
                self.garbage_collect()
                
                if self._get_total_allocated() + size > self.max_memory_bytes:
                    return None  # 内存不足
            
            offset = self.pools[pool_name].allocate(size)
            if offset is not None:
                obj_id = id((pool_name, offset))
                self.allocated_objects[obj_id] = pool_name
                return obj_id
            
            return None
    
    def deallocate(self, obj_id: int):
        """释放内存"""
        with self.lock:
            if obj_id not in self.allocated_objects:
                return
            
            pool_name = self.allocated_objects.pop(obj_id)
            # 这里需要从obj_id反推出offset，简化处理
            # 实际实现中需要维护更详细的映射关系
    
    def _select_pool_for_size(self, size: int) -> str:
        """根据大小选择合适的内存池"""
        if size <= 1024 * 1024:  # 1MB
            return 'small'
        elif size <= 16 * 1024 * 1024:  # 16MB
            return 'medium'
        else:
            return 'large'
    
    def _get_total_allocated(self) -> int:
        """获取总分配内存"""
        total = 0
        for pool in self.pools.values():
            stats = pool.get_usage_stats()
            total += stats['allocated_bytes']
        return total
    
    def garbage_collect(self):
        """执行垃圾回收"""
        gc.collect()
    
    def get_memory_stats(self) -> Dict:
        """获取内存统计信息"""
        with self.lock:
            stats = {
                'max_memory_bytes': self.max_memory_bytes,
                'total_allocated_bytes': self._get_total_allocated(),
                'system_memory': self._get_system_memory_info(),
                'pools': {}
            }
            
            for pool_name, pool in self.pools.items():
                stats['pools'][pool_name] = pool.get_usage_stats()
            
            return stats
    
    def _get_system_memory_info(self) -> Dict:
        """获取系统内存信息"""
        try:
            process = psutil.Process(os.getpid())
            memory_info = process.memory_info()
            
            return {
                'rss_bytes': memory_info.rss,  # 物理内存
                'vms_bytes': memory_info.vms,  # 虚拟内存
                'percent': process.memory_percent()
            }
        except:
            return {
                'rss_bytes': 0,
                'vms_bytes': 0,
                'percent': 0.0
            }
    
    def check_memory_pressure(self) -> bool:
        """检查内存压力"""
        stats = self.get_memory_stats()
        utilization = stats['total_allocated_bytes'] / stats['max_memory_bytes']
        
        return utilization > 0.8  # 80%阈值
    
    def cleanup(self):
        """清理资源"""
        with self.lock:
            self.allocated_objects.clear()
            for pool in self.pools.values():
                pool.free_blocks.clear()
                pool.allocated_blocks.clear()
                pool.allocated_size = 0

# 全局内存管理器实例
_global_memory_manager = None

def get_memory_manager() -> MemoryManager:
    """获取全局内存管理器实例"""
    global _global_memory_manager
    if _global_memory_manager is None:
        _global_memory_manager = MemoryManager()
    return _global_memory_manager

def set_memory_manager(manager: MemoryManager):
    """设置全局内存管理器实例"""
    global _global_memory_manager
    _global_memory_manager = manager

